import random
import logging
from typing import List, Dict, Set
import pandas as pd

from source.RFT.utils import logger, reconstruct_traces, trace_contains_component
from source.data_processor.utils.experiment_data import ExperimentData

def filter_and_sample_traces(
    exp_data: ExperimentData,
    min_size: int,
    max_samples: int,
    anomalous_trace_ids: Set[str]
) -> List[str]:
    """
    筛选和采样Trace
    
    参数:
        exp_data: 实验数据对象
        min_size: Trace最小长度
        max_samples: 每个场景最大采样数量
        anomalous_trace_ids: 异常Trace ID集合
        
    返回:
        采样后的Trace ID列表
    """
    logger.info(f"开始筛选和采样Trace: anomaly_component={exp_data.anomaly_component}, min_size={min_size}, max_samples={max_samples}")
    
    # 检查依赖关系数据是否存在
    if exp_data.dependencies_df is None or exp_data.dependencies_df.empty:
        logger.warning(f"依赖关系数据为空，无法筛选Trace")
        return []
    
    # 确保必要的列存在
    required_cols = ["traceId", "cmdb_id"]
    if not all(col in exp_data.dependencies_df.columns for col in required_cols):
        logger.warning(f"依赖关系DataFrame缺少必要列: {required_cols}")
        return []
    
    # 1. 直接从dependencies_df中筛选出异常Trace
    df = exp_data.dependencies_df
    
    # 将anomalous_trace_ids转换为列表，以便用于pandas的isin方法
    anomalous_list = list(anomalous_trace_ids)
    
    # 筛选出traceId在anomalous_trace_ids中的数据
    anomalous_df = df[df["traceId"].isin(anomalous_list)]
    
    if anomalous_df.empty:
        logger.warning(f"没有找到异常Trace数据")
        return []
    
    # 2. 按traceId分组，筛选出满足条件的Trace
    valid_trace_ids = []
    skipped_reason = {"no_anomaly_component": 0, "too_short": 0}
    
    # 获取所有唯一的traceId
    unique_trace_ids = anomalous_df["traceId"].unique()
    logger.info(f"找到 {len(unique_trace_ids)} 个异常Trace")
    
    for trace_id in unique_trace_ids:
        # 获取当前trace_id的所有数据
        trace_df = anomalous_df[anomalous_df["traceId"] == trace_id]
        
        # 条件1: 必须包含故障组件
        if exp_data.anomaly_component not in trace_df["cmdb_id"].values:
            skipped_reason["no_anomaly_component"] += 1
            continue
        
        # 条件2: 必须满足最小长度
        if len(trace_df) < min_size:
            skipped_reason["too_short"] += 1
            continue
        
        valid_trace_ids.append(trace_id)
    
    logger.info(f"筛选结果: {len(valid_trace_ids)}/{len(unique_trace_ids)} 个Trace通过筛选")
    logger.debug(f"筛选细节: 不包含故障组件={skipped_reason['no_anomaly_component']}, 长度过短={skipped_reason['too_short']}")
    
    # 3. 随机采样
    if not valid_trace_ids:
        logger.warning(f"没有符合条件的Trace")
        return []
    
    sample_count = min(len(valid_trace_ids), max_samples)
    sampled_traces = random.sample(valid_trace_ids, sample_count)
    
    logger.info(f"采样完成: 从 {len(valid_trace_ids)} 个有效Trace中采样 {len(sampled_traces)} 个")
    
    return sampled_traces

